#! /usr/bin/python3

import argparse
from multiprocessing import Pool, Value
import time
import random


def argument():
    arg = argparse.ArgumentParser()
    arg.add_argument("--file", '-f', required = True, help = 'Fasta file name')
    arg.add_argument("--threads", '-t', default = 1, type = int, help = 'Number of threads')
    arg.add_argument("--warning", '-w', default = False, type = bool, help = 'Print Warning information')
    arg.add_argument("--needmemsave", '-m', default = False, type = bool, help = 'Need memory save module, recommended >10000 sequences')
    return arg.parse_args()

sequences = []
__Warning__ = False

# compute the SP of two seqs
# by cjt
def Compute_two(place1: int, place2: int, d=0, e=0, m=1, mis=0) -> float:
    s1 = sequences[place1]
    s2 = sequences[place2]
    s1 = s1.replace(".", "-")
    s2 = s2.replace(".", "-")
    if len(s1) != len(s2):
        # print(s1, s2)
        if __Warning__: print("Warning: the length of s1 and s2 is wrong!")
    maxlen = max(len(s1), len(s2))
    if len(s1) != maxlen: s1 = s1 + '-' * (maxlen - len(s1))
    if len(s2) != maxlen: s2 = s2 + '-' * (maxlen - len(s2))
    score_two = 0
    gap1 = 0
    for i in range(maxlen):
        if s1[i] != "-" and s2[i] != "-":
            if gap1:
               gap1 = 0
            if s1[i] == s2[i]:
                score_two += m
            else:
                score_two += mis
        elif s1[i] != "-" or s2[i] != "-":
            if gap1 == 0:
                score_two -= d
                gap1 = 1
            else:
                score_two -= e
    #print("OK: ", place1, place2)
    return score_two

def Calc_SP_Scores_memsave(threads):
    Value_SP = 0
    Value_All = 0
    len_sequences = len(sequences)
    for i in range(len_sequences):
        result = []
        pool = Pool(processes = threads)
        leni = len(sequences[i])
        for j in range(i + 1, len_sequences):
            result.append( pool.apply_async( Compute_two, (i, j, ) ) )
            Value_All += leni
        pool.close()
        pool.join()
        for data in result:
            Value_SP += data.get()
        del pool
        print("\r STEP {0:} / {1:} ".format(i, len_sequences), end = '')
    print("OK")
    return Value_SP / Value_All

def Calc_SP_Scores(threads):
    Value_SP = 0
    Value_All = 0
    len_sequences = len(sequences)
    result = []
    pool = Pool(processes = threads)
    for i in range(len_sequences):
        leni = len(sequences[i])
        for j in range(i + 1, len_sequences):
            result.append( pool.apply_async( Compute_two, (i, j, ) ) )
            Value_All += leni
        print("\r STEP {0:} / {1:} ".format(i, len_sequences), end = '')
    pool.close()
    pool.join()
    for data in result:
        Value_SP += data.get()
    del pool
    print("OK")
    return Value_SP / Value_All

if __name__ == '__main__':
    start_real_time = time.time()
    start_CPU_time = time.perf_counter()
    arg = argument()
    f = open(arg.file)
    needmemsave = arg.needmemsave
    threads = arg.threads
    __Warning__ = arg.warning
    thisline = []
    for line in f:
        if line.startswith(">"):
            if thisline != []: sequences.append(''.join(thisline))
            thisline = []
        else:
            thisline.append(line.strip())
    if thisline != []:
        sequences.append(''.join(thisline))
    # print(sequences)
    if len(sequences) < 1: 
        print("Error: this file has no sequences!!")
        exit(1)
    # Check the sequence is the same length
    seqlen1 = len(sequences[0])
    seqlen = [seqlen1]
    maxlen = seqlen1
    sequence_cnt = 1
    for i in range(1, len(sequences)):
        thislen = len(sequences[i])
        if thislen != seqlen1:
            if __Warning__: print("Warning: The length of sequence " + str(i) + "(" + str(thislen) + ") is not same as sequence 0(" + str(seqlen1) + ")!!")
            maxlen = max(maxlen, thislen)
        seqlen.append(thislen)
        sequence_cnt += 1
    if needmemsave: Value_SP = Calc_SP_Scores_memsave(threads)
    else: Value_SP = Calc_SP_Scores(threads)
    end_real_time = time.time()
    end_CPU_time = time.perf_counter()
    print("SP Score is: " + str(Value_SP) + "\nUsed time: " + str(end_real_time - start_real_time) + ", CPU time: " + str(end_CPU_time - start_CPU_time))
